// prefetcher.c
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/resource.h>
#include <bpf/libbpf.h>
#include "prefetcher.h"
#include "prefetcher.skel.h" // 由 libbpf 自动生成

#define PREFETCH_PAGES 4 // 对应你设计中的预测未来P个，这里P=4
#define PAGE_SIZE 4096

static volatile bool exiting = false;

// 信号处理函数，用于优雅退出
static void sig_handler(int sig) {
    exiting = true;
}

// 核心函数：虚拟地址到文件偏移的转换
static int vma_to_file_offset(pid_t pid, unsigned long vma, char *filepath, size_t filepath_len, off_t *offset) {
    char maps_path[64];
    snprintf(maps_path, sizeof(maps_path), "/proc/%d/maps", pid);

    FILE *f = fopen(maps_path, "r");
    if (!f) return -1;

    char *line = NULL;
    size_t len = 0;
    int found = -1;

    while (getline(&line, &len, f) != -1) {
        unsigned long start, end, file_off;
        char perms[5], dev[6];
        long inode;
        char current_path[4096] = {0};

        sscanf(line, "%lx-%lx %4s %lx %5s %ld %s", &start, &end, perms, &file_off, dev, &inode, current_path);
        
        // 确保是文件映射并且可读
        if (vma >= start && vma < end && inode != 0 && strchr(perms, 'r')) {
            strncpy(filepath, current_path, filepath_len - 1);
            filepath[filepath_len - 1] = '\0';
            *offset = file_off + (vma - start);
            found = 0;
            break;
        }
    }

    free(line);
    fclose(f);
    return found;
}


// Ring Buffer 事件处理回调
static int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct event *e = data;
    printf("--- Potential Pattern Detected ---\n");
    printf("PID: %u, Fault Address: 0x%lx, Stride: %ld bytes (%ld KB)\n", 
           e->pid, e->fault_addr, e->stride, e->stride / 1024);

    // 1. 确认模式 (在高级版本中加入状态机和置信度)
    // 目前简化处理，直接执行预取

    // 2. 预测地址
    unsigned long next_addr = e->fault_addr + e->stride;
    printf("Predicting next access at: 0x%lx\n", next_addr);

    // 3. 虚拟地址 -> 文件偏移
    char filepath[4096];
    off_t file_offset;
    if (vma_to_file_offset(e->pid, next_addr, filepath, sizeof(filepath), &file_offset) != 0) {
        fprintf(stderr, "WARN: Could not translate VMA 0x%lx for PID %d to file offset.\n", next_addr, e->pid);
        return 0;
    }
    
    // 4. 执行预取
    int fd = open(filepath, O_RDONLY);
    if (fd < 0) {
        perror("WARN: Could not open file for prefetching");
        return 0;
    }

    off_t prefetch_offset = file_offset;
    size_t prefetch_size = PREFETCH_PAGES * PAGE_SIZE;

    printf("Prefetching: file=%s, offset=%ld, size=%zu\n", filepath, prefetch_offset, prefetch_size);

    // 使用 posix_fadvise 执行预取
    if (posix_fadvise(fd, prefetch_offset, prefetch_size, POSIX_FADV_WILLNEED) != 0) {
        perror("WARN: posix_fadvise failed");
    }

    close(fd);
    return 0;
}


int main(int argc, char **argv) {
    struct prefetcher_bpf *skel;
    struct ring_buffer *rb = NULL;
    int err;

    // 提升资源限制，对于加载 eBPF 程序是必要的
    struct rlimit r = {RLIM_INFINITY, RLIM_INFINITY};
    if (setrlimit(RLIMIT_MEMLOCK, &r)) {
        perror("setrlimit(RLIMIT_MEMLOCK)");
        return 1;
    }

    // 加载并验证 eBPF 程序
    skel = prefetcher_bpf__open_and_load();
    if (!skel) {
        fprintf(stderr, "Failed to open and load BPF skeleton\n");
        return 1;
    }

    // 挂载 tracepoint
    err = prefetcher_bpf__attach(skel);
    if (err) {
        fprintf(stderr, "Failed to attach BPF skeleton\n");
        goto cleanup;
    }

    printf("eBPF prefetcher started. Press Ctrl+C to exit.\n");

    // 设置 Ring Buffer 回调
    rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event, NULL, NULL);
    if (!rb) {
        err = -1;
        fprintf(stderr, "Failed to create ring buffer\n");
        goto cleanup;
    }
    
    // 设置信号处理
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    // 轮询 Ring Buffer
    while (!exiting) {
        err = ring_buffer__poll(rb, 100 /* 100ms timeout */);
        if (err == -EINTR) {
            err = 0;
            break;
        }
        if (err < 0) {
            printf("Error polling ring buffer: %d\n", err);
            break;
        }
    }

cleanup:
    printf("\nDetaching and cleaning up...\n");
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);
    return -err;
}